#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <eigen3/Eigen/Dense>
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>
#include <cmath>
#include <algorithm>

using namespace std::chrono_literals;
static const rclcpp::Logger LOGGER = rclcpp::get_logger("force_servo_control_node_adrc");

class ForceServoControlNode : public rclcpp::Node
{
public:
  ForceServoControlNode() : Node("force_servo_control_node")
  {
    // —— 主流程参数 —— //
    vx_ = declare_parameter<double>("vx", 0.02);                   // x 向恒速 m/s
    stop_x_threshold_ = declare_parameter<double>("stop_x", 0.1);  // 到此停止 m
    publish_period_ms_ = declare_parameter<int>("period_ms", 40);  // 25 Hz
    cmd_topic_ = declare_parameter<std::string>("cmd_topic", "force_servo_control_node/delta_twist_cmds");

    // —— 目标点（一次性）—— //
    go_to_target_ = declare_parameter<bool>("go_to_target", true);
    target_x_ = declare_parameter<double>("target_x", 0.05);  // m（base_link）
    target_z_ = declare_parameter<double>("target_z", 0.20);  // m
    pos_k_    = declare_parameter<double>("pos_k", 1.0);      // 位置 P 增益（m/s per m）
    pos_vmax_ = declare_parameter<double>("pos_vmax", 0.05);  // 目标阶段速度限幅 m/s
    pos_tol_  = declare_parameter<double>("pos_tol", 0.003);  // 到达判定阈值 m

    // —— 力控（z 轴 ADRC 恒力）—— //
    // F_low_/F_high_仅用于约束F_ref_与监控；力控完全由ADRC实现
    F_low_   = declare_parameter<double>("F_low", 0.0);
    F_high_  = declare_parameter<double>("F_high", 1.0);
    v_z_max_ = declare_parameter<double>("v_z_max", 0.02);   // z向速度限幅 m/s
    down_sign_ = declare_parameter<double>("down_sign", -1.0); // 向下方向（+z向上则取-1）
    F_ref_   = declare_parameter<double>("F_ref", 0.6);     // 期望恒力 (N)

    // —— ADRC 参数 —— //
    adrc_.b0     = declare_parameter<double>("adrc.b0", 50.0);
    adrc_.r0     = declare_parameter<double>("adrc.r0", 80.0);
    adrc_.beta01 = declare_parameter<double>("adrc.beta01", 200.0);
    adrc_.beta02 = declare_parameter<double>("adrc.beta02", 3000.0);
    adrc_.beta03 = declare_parameter<double>("adrc.beta03", 20000.0);
    adrc_.delta  = declare_parameter<double>("adrc.delta", 0.02);
    adrc_.k1     = declare_parameter<double>("adrc.k1", 2.0);
    adrc_.k2     = declare_parameter<double>("adrc.k2", 0.2);
    adrc_.alpha1 = declare_parameter<double>("adrc.alpha1", 0.8);
    adrc_.alpha2 = declare_parameter<double>("adrc.alpha2", 1.2);

    // 订阅力数据（只用 z，建议外部已做零点/重力补偿）
    force_sub_ = create_subscription<geometry_msgs::msg::TwistStamped>(
        "/force_data", 10,
        std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

    // TF
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

    // 发布给 Servo 的 Twist
    twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>(cmd_topic_, 10);
  }

  void initialize()
  {
    setupPlanningSceneMonitor();
    setupServo();

    state_ = go_to_target_ ? State::GO_TO_TARGET : State::RUN;

    timer_ = create_wall_timer(
        std::chrono::milliseconds(publish_period_ms_),
        std::bind(&ForceServoControlNode::control_loop, this));

    RCLCPP_INFO(LOGGER, "ADRC力控版本启动：vx=%.3f, stop_x=%.3f, v_z_max=%.3f, down_sign=%.1f",
                vx_, stop_x_threshold_, v_z_max_, down_sign_);
    RCLCPP_INFO(LOGGER, "go_to_target=%s, target_x=%.3f, target_z=%.3f, pos_k=%.3f, pos_vmax=%.3f, pos_tol=%.3f",
                go_to_target_ ? "true":"false", target_x_, target_z_, pos_k_, pos_vmax_, pos_tol_);
    RCLCPP_INFO(LOGGER, "F_ref=%.3f (clamped to [%.2f, %.2f]), adrc: b0=%.1f r0=%.1f beta=[%.1f,%.1f,%.1f] delta=%.3f k=[%.2f,%.2f] alpha=[%.2f,%.2f]",
                F_ref_, F_low_, F_high_, adrc_.b0, adrc_.r0, adrc_.beta01, adrc_.beta02, adrc_.beta03,
                adrc_.delta, adrc_.k1, adrc_.k2, adrc_.alpha1, adrc_.alpha2);
  }

private:
  // —— 通信对象 —— //
  rclcpp::Subscription<geometry_msgs::msg::TwistStamped>::SharedPtr force_sub_;
  rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
  std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
  std::shared_ptr<planning_scene_monitor::PlanningSceneMonitor> planning_scene_monitor_;
  std::unique_ptr<moveit_servo::Servo> servo_;
  rclcpp::TimerBase::SharedPtr timer_;

  // —— 状态 & 变量 —— //
  enum class State { GO_TO_TARGET, RUN, STOPPED };
  State state_{State::RUN};
  Eigen::Vector3d current_force_{Eigen::Vector3d::Zero()};

  // —— 参数 —— //
  double vx_;
  double stop_x_threshold_;
  int publish_period_ms_;
  std::string cmd_topic_;

  // 目标点
  bool   go_to_target_;
  double target_x_, target_z_, pos_k_, pos_vmax_, pos_tol_;

  // 力控（ADRC）
  double F_low_, F_high_;   // 期望力边界（用于夹持/监控）
  double v_z_max_;          // z向速度限幅
  double down_sign_;        // z向“向下”符号
  double F_ref_;            // 期望恒力
  double last_u_z_{0.0};    // ESO用：上一周期控制输入

  // ========= ADRC 实现 =========
  static inline double sgn(double x) { return (x > 0) - (x < 0); }

  static double fal(double e, double alpha, double delta)
  {
    double ae = std::abs(e);
    if (ae <= delta) return e / std::pow(delta, 1.0 - alpha);
    return std::pow(ae, alpha) * sgn(e);
  }

  static double fhan(double x1, double x2, double r, double h)
  {
    double d  = r * h;
    double d0 = h * d;
    double y  = x1 + h * x2;
    double a0 = std::sqrt(d * d + 8 * r * std::abs(y));
    double a;
    if (std::abs(y) > d0) a = x2 + (a0 - d) * 0.5 * sgn(y);
    else                  a = x2 + y / h;

    if (std::abs(a) > d)  return -r * sgn(a);
    return -r * a / d;
  }

  struct ADRCz
  {
    // 参数（可从参数服务器配置）
    double b0{50.0};   // 模型增益（把控制量映射到力/速度域的经验参数）
    double r0{80.0};   // TD 速度系数（小=更柔和）
    double beta01{200.0}, beta02{3000.0}, beta03{20000.0}; // ESO增益
    double delta{0.02};  // fal小区间
    double k1{2.0}, k2{0.2};       // 非线性组合增益
    double alpha1{0.8}, alpha2{1.2};

    // 状态
    double v1{0.0}, v2{0.0};           // TD：期望力及其导数
    double z1{0.0}, z2{0.0}, z3{0.0};  // ESO：测力、导数、总扰动

    // 单步：输入Fr(期望力), Fe(测量力), u_prev(上一周期控制输入), h(周期)
    double step(double Fr, double Fe, double u_prev, double h)
    {
      // 1) TD：平滑把设定拉向Fr
      double f0 = ForceServoControlNode::fhan(v1 - Fr, v2, r0, h);
      v1 += h * v2;
      v2 += h * f0;

      // 2) ESO：估计z1,z2,z3（用fal抑噪）
      double e   = z1 - Fe;
      double fe  = ForceServoControlNode::fal(e, 0.5, delta);
      double fe1 = ForceServoControlNode::fal(e, 0.5, delta);
      z1 += h * ( z2 - beta01 * e );
      z2 += h * ( z3 - beta02 * fe + u_prev );
      z3 += h * ( -beta03 * fe1 );

      // 3) 非线性组合 + 扰动补偿
      double e1 = v1 - z1;
      double e2 = v2 - z2;
      double u0 = k1 * ForceServoControlNode::fal(e1, alpha1, delta)
                + k2 * ForceServoControlNode::fal(e2, alpha2, delta);
      return (u0 - z3) / b0; // 控制量u
    }
  } adrc_;

  // ———— 回调：力数据（只用 z）———— //
  void force_callback(const geometry_msgs::msg::TwistStamped::SharedPtr msg)
  {
    current_force_.x() = 0.0;
    current_force_.y() = 0.0;
    current_force_.z() = msg->twist.linear.z;  // 单位 N（建议外部已做零点/重力补偿）
  }

  // ———— 主循环 ———— //
  void control_loop()
  {
    // 读取当前末端在 base_link 下的位置
    double cur_x = 0.0, cur_z = 0.0;
    try {
      auto tf = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
      cur_x = tf.transform.translation.x;
      cur_z = tf.transform.translation.z;
    } catch (tf2::TransformException &ex) {
      RCLCPP_WARN(LOGGER, "获取末端位置失败：%s", ex.what());
      return;
    }

    auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
    msg->header.stamp = this->now();
    msg->header.frame_id = "base_link";

    if (state_ == State::GO_TO_TARGET)
    {
      // 只在 x、z 两个方向靠近目标（y 不动）
      double ex = target_x_ - cur_x;
      double ez = target_z_ - cur_z;

      double vx_cmd = clamp(pos_k_ * ex, -pos_vmax_, pos_vmax_);
      double vz_cmd = clamp(pos_k_ * ez, -pos_vmax_, pos_vmax_);

      // 到达判据（x、z 都进 tol）
      if (std::abs(ex) < pos_tol_ && std::abs(ez) < pos_tol_)
      {
        vx_cmd = 0.0; vz_cmd = 0.0;
        state_ = State::RUN;
        RCLCPP_INFO(LOGGER, "已到达目标 (x=%.3f, z=%.3f)，切换到 RUN", cur_x, cur_z);
      }

      msg->twist.linear.x = vx_cmd;
      msg->twist.linear.y = 0.0;
      msg->twist.linear.z = vz_cmd;
      msg->twist.angular.x = 0.0;
      msg->twist.angular.y = 0.0;
      msg->twist.angular.z = 0.0;

      twist_cmd_pub_->publish(std::move(msg));
      return;
    }

    if (state_ == State::RUN)
    {
      // 到达 x 阈值就全停
      if (cur_x >= stop_x_threshold_) {
        state_ = State::STOPPED;
        RCLCPP_INFO(LOGGER, "末端 x=%.3f >= %.3f，进入 STOPPED", cur_x, stop_x_threshold_);
      }
    }

    if (state_ == State::STOPPED)
    {
      // 全停
      msg->twist.linear.x = 0.0;
      msg->twist.linear.y = 0.0;
      msg->twist.linear.z = 0.0;
      msg->twist.angular.x = 0.0;
      msg->twist.angular.y = 0.0;
      msg->twist.angular.z = 0.0;
      twist_cmd_pub_->publish(std::move(msg));
      return;
    }

    // —— RUN 状态：x 恒速 + z 轴 ADRC 恒力 —— //
    const double Fe = current_force_.z();                   // 测量力
    const double Fr = clamp(F_ref_, F_low_, F_high_);       // 期望力（约束在安全区）
    const double h  = publish_period_ms_ * 1e-3;            // 控制周期(s)

    // ADRC 一步（u为“速度指令规模”）：用上一周期u作为ESO输入
    const double u = adrc_.step(Fr, Fe, last_u_z_, h);
    last_u_z_ = u;

    // 按down_sign_映射到z向速度（正负取决于坐标系约定）
    double v_z_cmd = down_sign_ * u;

    // 限幅
    if (v_z_cmd >  v_z_max_) v_z_cmd =  v_z_max_;
    if (v_z_cmd < -v_z_max_) v_z_cmd = -v_z_max_;

    // 发布速度
    msg->twist.linear.x = vx_;       // 沿 x 恒速
    msg->twist.linear.y = 0.0;
    msg->twist.linear.z = v_z_cmd;   // z 轴 ADRC 输出
    msg->twist.angular.x = 0.0;
    msg->twist.angular.y = 0.0;
    msg->twist.angular.z = 0.0;

    twist_cmd_pub_->publish(std::move(msg));
  }

  // —— 工具：限幅 —— //
  static double clamp(double v, double lo, double hi)
  { return std::max(lo, std::min(v, hi)); }

  // —— PSM & Servo —— //
  void setupPlanningSceneMonitor()
  {
    planning_scene_monitor_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
        shared_from_this(), "robot_description", tf_buffer_, "planning_scene_monitor");

    if (planning_scene_monitor_->getPlanningScene())
    {
      planning_scene_monitor_->startStateMonitor("/joint_states");
      planning_scene_monitor_->setPlanningScenePublishingFrequency(25);
      planning_scene_monitor_->startPublishingPlanningScene(
          planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
          "/moveit_servo/publish_planning_scene");
      planning_scene_monitor_->startSceneMonitor();
      planning_scene_monitor_->providePlanningSceneService();
    }
    else
    {
      RCLCPP_FATAL(LOGGER, "PlanningSceneMonitor 配置失败");
      throw std::runtime_error("Planning scene setup failed");
    }
  }

  void setupServo()
  {
    auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(shared_from_this());
    if (!servo_parameters)
    {
      RCLCPP_FATAL(LOGGER, "无法加载 MoveIt Servo 参数");
      throw std::runtime_error("Failed to load servo parameters");
    }

    servo_ = std::make_unique<moveit_servo::Servo>(shared_from_this(), servo_parameters, planning_scene_monitor_);
    servo_->start();
    RCLCPP_INFO(LOGGER, "MoveIt Servo 已启动，监听话题：%s", cmd_topic_.c_str());
  }
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ForceServoControlNode>();
  node->initialize();   // 构造后再初始化，避免 bad_weak_ptr
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
